"""
FunASR模型加载调试脚本

用途：测试FunASR模型加载是否正常，并分析可能的问题
"""

import os
import sys
import torch
import importlib
import traceback

# 获取项目根目录
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

# 设置模型路径
MODEL_DIR = os.path.join(ROOT_DIR, "models", "sensevoice", "SenseVoiceSmall")
FSMN_VAD_DIR = os.path.join(ROOT_DIR, "models", "fsmn", "speech_fsmn_vad_zh-cn-16k-common-pytorch")
MODEL_PY = os.path.join(ROOT_DIR, "model.py")

def check_paths():
    """检查路径是否存在"""
    print("检查路径...")
    print(f"当前目录: {os.getcwd()}")
    print(f"脚本目录: {ROOT_DIR}")
    
    paths_to_check = [
        ("模型目录", MODEL_DIR),
        ("VAD模型目录", FSMN_VAD_DIR),
        ("模型代码文件", MODEL_PY)
    ]
    
    all_exist = True
    for name, path in paths_to_check:
        exists = os.path.exists(path)
        print(f"{name}: {path} {'存在' if exists else '不存在'}")
        if not exists:
            all_exist = False
    
    return all_exist

def check_funasr_import():
    """测试导入FunASR包"""
    print("\n测试导入FunASR...")
    try:
        import funasr
        print(f"FunASR版本: {funasr.__version__ if hasattr(funasr, '__version__') else '未知'}")
        
        # 检查自动模型类
        try:
            from funasr import AutoModel
            print("FunASR AutoModel 导入成功")
        except ImportError as e:
            print(f"无法导入 AutoModel: {e}")
            return False
            
        # 检查其他关键模块
        modules_to_check = [
            "funasr.utils",
            "funasr.models",
            "funasr.register"
        ]
        
        for module_name in modules_to_check:
            try:
                module = importlib.import_module(module_name)
                print(f"导入 {module_name} 成功")
            except ImportError as e:
                print(f"无法导入 {module_name}: {e}")
                return False
        
        return True
    except ImportError as e:
        print(f"无法导入 FunASR: {e}")
        return False
    except Exception as e:
        print(f"导入 FunASR 时发生未知错误: {e}")
        print(traceback.format_exc())
        return False

def test_model_load():
    """测试加载模型"""
    print("\n测试加载模型...")
    
    # 加载AutoModel
    try:
        from funasr import AutoModel
        print("步骤1: 导入AutoModel完成")
        
        # 检查各个路径
        for path_name, path in [
            ("模型路径", MODEL_DIR),
            ("VAD模型路径", FSMN_VAD_DIR),
            ("远程代码路径", MODEL_PY)
        ]:
            if not os.path.exists(path):
                print(f"错误: {path_name} 不存在: {path}")
                return False
        
        print("步骤2: 所有路径检查通过")
        
        # 检查torch可用性
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        print(f"使用设备: {device}")
        
        # 打印调试信息
        print(f"开始初始化模型，参数如下:")
        print(f"  - model: {MODEL_DIR}")
        print(f"  - trust_remote_code: True")
        print(f"  - remote_code: {MODEL_PY}")
        print(f"  - vad_model: {FSMN_VAD_DIR}")
        print(f"  - device: {device}")
        
        # 尝试加载模型
        model = AutoModel(
            model=MODEL_DIR,
            trust_remote_code=True,
            remote_code=MODEL_PY,
            vad_model=FSMN_VAD_DIR,
            vad_kwargs={"max_single_segment_time": 30000},
            device=device,
        )
        
        print("步骤3: 模型加载成功!")
        print(f"模型类型: {type(model)}")
        
        return True
    except Exception as e:
        print(f"初始化模型失败: {e}")
        print(traceback.format_exc())
        return False

def suggest_solutions():
    """提供可能的解决方案"""
    print("\n可能的解决方案:")
    print("1. 确保在打包前安装了正确版本的FunASR (>= 1.1.3)")
    print("2. 添加以下参数到build.py中的PyInstaller命令:")
    print("   --hidden-import=funasr.auto.auto_model")
    print("   --hidden-import=funasr.register")
    print("   --hidden-import=funasr.models")
    print("3. 将model.py文件放在打包目录的根目录中，确保它被正确包含")
    print("4. 尝试添加--collect-all=funasr参数来收集所有FunASR相关的模块")
    print("5. 检查模型文件是否完整，尝试在打包前重新下载模型")
    print("6. 如果问题仍然存在，尝试修改app.py中的AutoModel初始化部分，增加错误处理和回退机制")

def main():
    """主函数"""
    print("FunASR模型加载调试工具")
    print("====================\n")
    
    # 检查路径
    if not check_paths():
        print("\n错误: 一些必要的路径不存在，请确保模型和代码文件都在正确位置")
        suggest_solutions()
        return False
    
    # 检查FunASR导入
    if not check_funasr_import():
        print("\n错误: FunASR导入测试失败")
        suggest_solutions()
        return False
    
    # 测试模型加载
    if not test_model_load():
        print("\n错误: 模型加载测试失败")
        suggest_solutions()
        return False
    
    print("\n所有测试通过! FunASR模型可以正常加载。")
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
